#ifndef _HUFFMAN_TREE_H
#define _HUFFMAN_TREE_H

#include "bitstream.h"
#include <stdio.h>
#include <stdint.h>

template <class T> class cHuffmanItem;
class cHuffmanCode
{
public:
	cHuffmanCode();
	~cHuffmanCode();
	uint32_t * Bits;
	int16_t Length;
	int16_t Size;
	uint32_t bits;
	void Put(bool bit);
	void Write(cBitstream * bitstream);

};


enum cHuffmanItemType
{
	None = 0,
	Leaf = 1,
	Node = 2,
	New = 3,
	EOB = 4,
	ERROR
};
template <class T> 
class cHuffmanTree
{
public:
	cHuffmanTree();
	cHuffmanTree(const cHuffmanTree & source);
	~cHuffmanTree(void);
	bool EncodeBlock(void * ptr,uint32_t size,cBitstream * bstr= NULL,bool inc_new = true);
	enum cHuffmanItemType DecodeBlock(void * ptr,uint32_t size,cBitstream * bstr= NULL);
	bool Encode(T symbol,cBitstream * bitstream = NULL,bool inc_new = true);
	bool EncodeCtl(cHuffmanItemType itemtype,cBitstream * bitstream = NULL);
	enum cHuffmanItemType Decode(cBitstream * bitstream,T * symbol);
	void Reset();
	void Leak();
	uint32_t ItemCounter();
	uint32_t SymbolCounter();
	uint64_t Counter();
	void PrintAll();
	void DeleteAll();
	void setSymbolSize(uint8_t ss);
	uint8_t getSymbolSize();
  private:
	uint8_t symbol_size;
	uint32_t leafCounter;
	void Add(T value,uint32_t weight,cHuffmanItemType itemtype);
	cHuffmanItem<T> * First;
	cHuffmanItem<T> * Last;
	cHuffmanCode * GetCode(cHuffmanItem<T> * symbol);
	cHuffmanItem<T> * FindLeafValue(T Value);
	cHuffmanItem<T> * FindItemType(cHuffmanItemType itemtype);
	void Init();

	void Insert(T symbol,cHuffmanItemType itemtype = Leaf);
	void Copy(const cHuffmanTree & source);
};
template <class T> 
class cHuffmanItem
{
public:
	~cHuffmanItem();
	cHuffmanItem();
	cHuffmanItem(const cHuffmanItem & source);
  private:
	cHuffmanItem(T value,uint64_t weight,cHuffmanItem * next,cHuffmanItem * parent,cHuffmanItem * left,cHuffmanItem * right,cHuffmanItemType itemtype,cHuffmanItem * prev = NULL,uint32_t n = 0);
	void IncWeightAndModify();
	void DecWeight();
	void DecWeightAndModify();
	void Print();
	void Exchange(cHuffmanItem * p,bool update_children = true);
	cHuffmanItem * FindN(uint32_t n);
	T Value;
	uint32_t SN;
	uint64_t Weight;
	cHuffmanItem * Next;
	cHuffmanItem * Prev;	
	cHuffmanItem * Left;		
	cHuffmanItem * Right;	
	cHuffmanItem * Parent;		
	cHuffmanItemType ItemType;	

friend class cHuffmanTree<T>;
};



//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanTree<T>::cHuffmanTree() 
{
  this->symbol_size = 8*sizeof(T);
  this->Init();
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanTree<T>::cHuffmanTree(const cHuffmanTree<T> & source)
{
	this->symbol_size = source.symbol_size;
	this->First = NULL;
	this->Last = NULL;
	this->Copy(source);
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanTree<T>::~cHuffmanTree(void) {this->DeleteAll();}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T>
void cHuffmanTree<T>::setSymbolSize(uint8_t ss)
{
  this->symbol_size = ss;
}
template <class T>
uint8_t cHuffmanTree<T>::getSymbolSize()
{
  return this->symbol_size;
}

template <class T> 
void cHuffmanTree<T>::Copy(const cHuffmanTree<T> & source)
{
	if(this->First || this->Last)
		this->DeleteAll();
	this->Last = new cHuffmanItem<T>(*source.Last);
	this->leafCounter = source.leafCounter;
	cHuffmanItem<T> * p = this->Last;
	uint32_t n = 1;
	while(n < source.First->SN)
	{
		if(p->Prev == NULL)
		{
			p->Prev = this->Last->FindN(n+1);
			p->Prev->Next = p;
			
		}
		p=p->Prev;
		n++;
		
	}
	this->First = p;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanTree<T>::Init()				//Na poczatku jeden wezel i dwa liscie EOB i New
{
	cHuffmanItem<T> * iEOB = new cHuffmanItem<T>(0,1,NULL,NULL,NULL,NULL,EOB,NULL,3);
	cHuffmanItem<T> * iNew = new cHuffmanItem<T>(0,1,NULL,NULL,NULL,NULL,New,NULL,2);
	cHuffmanItem<T> * iNode = new cHuffmanItem<T>(0,2,NULL,NULL,iEOB,iNew,Node,iNew,1);
	iEOB->Next = iNew;
	iEOB->Parent = iNode;
	iNew->Next = iNode;
	iNew->Parent = iNode;
	iNew->Prev = iEOB;
	First = iEOB;
	Last = iNode;
	this->leafCounter = 0;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanTree<T>::DeleteAll()
{
	cHuffmanItem<T> * p = NULL;
	cHuffmanItem<T> * next = NULL;
	for(p=First;p;p=next)
	{
		next = p->Next;
		delete p;
	}
	Last = NULL;
	First = NULL;
	this->leafCounter = 0;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanTree<T>::Reset()
{
	this->DeleteAll();
	this->Init();
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanTree<T>::Insert(T symbol,cHuffmanItemType itemtype)		// Wstawianie nowego symbolu do drzewa
{
	this->leafCounter++;
	cHuffmanItem<T> * u = new cHuffmanItem<T>(First->Value,First->Weight,First,First,First->Left,First->Right,First->ItemType,NULL,First->SN+1);
	cHuffmanItem<T> * v = new cHuffmanItem<T>(symbol,1,u,First,NULL,NULL,itemtype,NULL,First->SN+2);
	u->Prev = v;
	First->Value = 0;
	First->Left = v;
	First->Right = u;
	First->ItemType = Node;
	First->Weight = u->Weight;
	First->Prev = u;
	First = v;
	First->Next = u; 
	cHuffmanItem<T> * parent = First->Parent;
	if(parent)
		parent->IncWeightAndModify();
}
template <class T> 
bool cHuffmanTree<T>::EncodeCtl(cHuffmanItemType itemtype,cBitstream * bitstream)
{
  cHuffmanItem<T> * v = this->FindItemType(itemtype);
  if(v==NULL)
  {
    fprintf(stderr,"Encode CTL v=null\n");
    return false;
  }
  if(bitstream != NULL)
  {
    cHuffmanCode * code = this->GetCode(v);
    code->Write(bitstream);
    delete code;
  }
  v->IncWeightAndModify();
  return true;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
bool cHuffmanTree<T>::EncodeBlock(void * ptr,uint32_t size,cBitstream * bstr,bool inc_new)
{
  uint8_t * bptr = (uint8_t*)ptr;
  while(size--)
    this->Encode(*(bptr++),bstr,inc_new);
  return true;
}
template <class T> 
enum cHuffmanItemType cHuffmanTree<T>::DecodeBlock(void * ptr,uint32_t size,cBitstream * bstr)
{
  uint8_t * bptr = (uint8_t*)ptr;
  cHuffmanItemType itype;
  T symbol;
  while(size--)
  {
    itype = this->Decode(bstr,&symbol);
    if(itype!=Leaf && itype != New)
      return itype;
    *(bptr++) = (uint8_t)symbol;
  }
  return None;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
bool cHuffmanTree<T>::Encode(T symbol,cBitstream * bitstream,bool inc_new)	
{
	cHuffmanItem<T> * v = this->FindLeafValue(symbol);						//sprawdzamy czy w drzewie jest juz kodowana wartosc
	if(v == NULL)													//NULL oznacza ze w drzewie nie ma tej wertosci
	{
		cHuffmanItem<T> * newType = this->FindItemType(New);		//szukamy liscia z symbolem New
		if(bitstream != NULL)										//to jest po to, ze czasami funkcja Encode pracuje jako zwykle dodawanie do drzewa bez zapisu do strumienia
		{
			cHuffmanCode * code = this->GetCode(newType);			//pobieramy kod symbolu New
			code->Write(bitstream);
			delete code;			
			bitstream->putBits(this->symbol_size,(uint32_t)symbol);
		}
		if(inc_new)
		  newType->IncWeightAndModify();								//zwiekszamy wage symbolu New
		this->Insert(symbol);										//i wstawiamy nowy element do drzewa
		return true;												
	}
	if(bitstream != NULL)											//jezeli jest juz symbol i chcemy zapisac do strumienia..
	{
		cHuffmanCode * code = this->GetCode(v);								//...to pobieramy jego kod i zapisujemy
		code->Write(bitstream);
		delete code;
	}
	v->IncWeightAndModify();										//zwiekszamy wage symbolu 
	return true;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
enum cHuffmanItemType cHuffmanTree<T>::Decode(cBitstream * bitstream,T * ret)		
{
	cHuffmanItem<T> * p = Last;												//ostatni element listy jest korzeniem calego drzewa
	while( p!=NULL && p->ItemType == Node)				//wedrujemy po drzewie zgodnie z odbieranymi bitami ze strumienia
	{
		int8_t dir = (int8_t)bitstream->getBits(1);
		if(dir == 1)
			p = p->Right;
		else
			p = p->Left;
	}
	if(p == NULL)
	{
	  fprintf(stderr,"cHuffmanTree: Decode: p == NULL\n");
	  exit(1);
	  return ERROR;
	}
	if(p->ItemType == EOB)
	{
	  p->IncWeightAndModify();
	  *ret=0;
	  return EOB;
	}
	if(p->ItemType == New)								//jezeli odczytano symbol nowego elementu...
	{
		T symbol = (T)bitstream->getBits(this->symbol_size);			//to odczytujemy ten element zapisany na tylu bitach ile wskazuje zmienna s_length (domyslnie 16) ale do RLC gdy przesylamy dlugosc ciagu zapisujemy na 6 bitach
		this->Encode(symbol);										//tutaj po prostu dodajemy element do drzewa wykorzystujac f. Encode z domyslnym argumenetem bitstream = NULL
		*ret = symbol;
		return New;
	}
	T s = p->Value;												//jezeli odczytalismy jakis symbol to go zwracamy
	cHuffmanItemType itype= p->ItemType;
	this->Encode(p->Value);											//aktualizacja drzewa
	*ret = s;
	return itype;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanCode * cHuffmanTree<T>::GetCode(cHuffmanItem<T> * symbol)							//funkcja zwraca kod elementu przechodzac w gore drzewa za pomoca wskaznika Parent
{
	cHuffmanCode * code = new cHuffmanCode();
	cHuffmanItem<T> * p = symbol->Parent;
	while(p != NULL)
	{
		code->Put(symbol == p->Right);
		symbol = p;
		p=symbol->Parent;
	}
	return code;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T> * cHuffmanTree<T>::FindLeafValue(T value)				//funkcja szukajaca w drzewie danego symbolu, jezeli nie znaleziono zwraca NULL
{
	cHuffmanItem<T> * p  = this->Last->Prev;							//zaczynamy szukac od przedostatniego elementu listy bo ostatni jest korzen a elementy o najwiekszej wadze sa na koncu
	while(p && ((p->ItemType == Leaf && p->Value != value) || p->ItemType == EOB || p->ItemType == New || p->ItemType == Node ))
		p=p->Prev;
	if(p && p->Value == value && p->ItemType == Leaf)
		return p;
	return NULL;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T> * cHuffmanTree<T>::FindItemType(cHuffmanItemType itemtype)		//funkcja zwracajaca wskaznik symboli New lub EOB
{
	if(itemtype == EOB)									//gdy szukamy EOB to zaczynamy od poczatku bo EOB bedzie mial zawsze mala wage
	{
		cHuffmanItem<T> * p = this->First;
		while(p && p->ItemType != EOB)
			p=p->Next;
		if(p && p->ItemType==EOB)
			return p;
		return NULL;
	}
	if(itemtype == New)									//gdy szukamy New to zaczynamy od konca bo New moze miec wysoka wage
	{
		cHuffmanItem<T> * p = this->Last->Prev;
		while(p && p->ItemType != New)
			p=p->Prev;
		if(p && p->ItemType==New)
			return p;
		return NULL;
	}
	return NULL;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanTree<T>::PrintAll() {Last->Print();}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T> * cHuffmanItem<T>::FindN(uint32_t n)
{
	if(this->SN == n)
		return this;
	cHuffmanItem<T> * left = NULL;
	cHuffmanItem<T> * right = NULL;
	if(this->Right)
		right = this->Right->FindN(n);
	if(right)
		return right;
	if(this->Left)
		left = this->Left->FindN(n);
	if(left)
		return left;
	return NULL;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanTree<T>::Leak()
{
	cHuffmanItem<T> * p = this->First;
	cHuffmanItem<T> * next = NULL;
	for(p = this->First;p;p=next)
	{
		next = p->Next;
		if(p->ItemType == Leaf)
		{
			p->DecWeightAndModify();
			if(First->Weight == 0)
			{
				this->leafCounter--;
				cHuffmanItem<T> * f = First;
				First = First->Next->Next;
				First->Prev = NULL;
				f->Next->Exchange(f->Parent);
				delete f->Next;
				delete f;
				next=First;
				this->FindItemType(New)->DecWeightAndModify();
			}
		}
	}
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
uint32_t cHuffmanTree<T>::ItemCounter()
{
	return this->First->SN;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
uint32_t cHuffmanTree<T>::SymbolCounter()
{
	return this->leafCounter;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
uint64_t cHuffmanTree<T>::Counter()
{
	return this->Last->Weight;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T>::cHuffmanItem():Value(0),SN(0),Weight(1),Next(NULL),Prev(NULL),Left(NULL),Right(NULL),Parent(NULL),ItemType(None){}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T>::cHuffmanItem(const cHuffmanItem<T> & source)
{
	this->Value = source.Value;
	this->SN = source.SN;
	this->Weight = source.Weight;
	this->ItemType = source.ItemType;
	this->Next = NULL;
	this->Prev = NULL;
	if(source.ItemType == Node)
	{
		this->Left = new cHuffmanItem<T>(*source.Left);
		this->Right = new cHuffmanItem<T>(*source.Right);
		this->Left->Next = this->Right;
		this->Right->Prev = this->Left;
		this->Left->Parent = this;
		this->Right->Parent = this;
	}
	else
	{
		this->Left = NULL;
		this->Right = NULL;
	}
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T>::cHuffmanItem(T value,uint64_t weight,cHuffmanItem * next,cHuffmanItem * parent,cHuffmanItem * left,cHuffmanItem * right,cHuffmanItemType itemtype,cHuffmanItem * prev,uint32_t n)
	:Value(value),SN(n),Weight(weight),Next(next),Prev(prev),Left(left),Right(right),Parent(parent),ItemType(itemtype)
{
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
cHuffmanItem<T>::~cHuffmanItem()
{	
	/*if(this->ItemType == Leaf)
		printf("Usuwam: [Leaf][%d][%d][%d]\n",this->Value,this->SN,this->Weight);
	if(this->ItemType == Node)
		printf("Usuwam: [Node][%d][%d][%d]\n",this->Value,this->SN,this->Weight);
	if(this->ItemType == New)
		printf("Usuwam: [New][%d][%d][%d]\n",this->Value,this->SN,this->Weight);
	if(this->ItemType == EOB)
		printf("Usuwam: [EOB][%d][%d][%d]\n",this->Value,this->SN,this->Weight);*/
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanItem<T>::IncWeightAndModify()					//zwiekszanie wagi danego elementu i jego przodkow oraz modyfikacja drzewa zgodnie z zalozeniami drzewa Huffmana
{
	this->Weight++;
	cHuffmanItem<T> * p = this;
	while(p->Next && p->Next->Weight == (this->Weight-1) )
		p=p->Next;
	if( p != this && p->Weight == (this->Weight-1) )
		this->Exchange(p);
	if(p->Parent) 
		p->Parent->IncWeightAndModify();
	return;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanItem<T>::DecWeightAndModify()
{
	cHuffmanItem<T> * p = this;
	while(p->Prev && p->Prev->Weight == this->Weight) p=p->Prev;
	if(p != this && p->Weight == this->Weight)
	{
		this->Exchange(p);
	}
	p->Weight--;
	if(p->Parent)
		p->Parent->DecWeightAndModify();
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanItem<T>::DecWeight()
{
	this->Weight--;
	for(cHuffmanItem<T> * p = this->Parent; p ; p = p->Parent)
		p->Weight--;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanItem<T>::Exchange(cHuffmanItem<T> * p,bool update_children)			//zamiana dwoch wierzcholkow wraz z jego potomkami, Next, Prev i Parent sie nie zmienia!!!
{
	T value = p->Value;
	uint64_t weight = p->Weight;
	cHuffmanItem<T> * left = p->Left;
	cHuffmanItem<T> * right = p->Right;
	cHuffmanItemType itemtype = p->ItemType;
	p->Value = this->Value;
	p->Weight = this->Weight;
	p->Left = this->Left;
	p->Right = this->Right;
	p->ItemType = this->ItemType;
	this->Value = value;
	this->Weight = weight;
	this->Left = left;
	this->Right = right;
	this->ItemType = itemtype;
	if(!update_children)
		return;
	if(p->Left) 
		p->Left->Parent = p;
	if(p->Right)
		p->Right->Parent = p;
	if(this->Left)
		this->Left->Parent = this;
	if(this->Right)
		this->Right->Parent = this;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class T> 
void cHuffmanItem<T>::Print()
{
	if(Left)
		Left->Print();
	if(Right)
		Right->Print();
	if(this->ItemType == Node)
	{
		printf("[Node][%d]:Weight = %u val=%d \n",this->SN,(unsigned int)this->Weight,this->Value);
		return;
	}
	if(this->ItemType == New)
	{
		printf("[New][%d]:Weight = %u val=%d \n",this->SN,(unsigned int)this->Weight,this->Value);
		return;
	}
	if(this->ItemType == EOB)
	{
		printf("[EOB][%d]:Weight = %u val=%d\n",this->SN,(unsigned int)this->Weight,this->Value);
		return;
	}
	if(this->ItemType == Leaf)
	{
		printf("[%d][%d]:Weight = %u itype=%d \n",this->Value,this->SN,(unsigned int)this->Weight,this->ItemType);
		return;
	}
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#endif //_HUFFMAN_TREE_H